{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Example: linear evolution without solving the Grad-Shafranov equation\n", "\n", "This notebook follows on from the previous one, except that the following (linear) evolutive solves **will not** solve the GS equation at each timestep. \n", "\n", "This mode uses the linearisations from the initial equilibrium configuration at time $t=0$ (i.e. the Jacobians inside the evolutive solver) to evolve plasma shape parameters of interest (e.g. midplane radii, X-point positions) forward in time rapidly.\n", "\n", "This can be useful for doing very fast scans forward in time into how certain shape parameters will evolve (linearly at least). Note, however, the equilibria for these simulations cannot be visualised as GS has not been solved!\n", "\n", "The simulation proceeds in almost exactly the same way as the previous notebook so please do check that out first. " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Static forward simulation\n", "\n", "Begin by building the initial plasma configuration as usual. " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "import numpy as np\n", "\n", "# build machine\n", "from freegsnke import build_machine\n", "tokamak = build_machine.tokamak(\n", " active_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_active_coils.pickle\",\n", " passive_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_passive_coils.pickle\",\n", " limiter_path=f\"../machine_configs/MAST-U/MAST-U_like_limiter.pickle\",\n", " wall_path=f\"../machine_configs/MAST-U/MAST-U_like_wall.pickle\",\n", ")\n", "\n", "# initialise equilibrium object\n", "from freegsnke import equilibrium_update\n", "eq = equilibrium_update.Equilibrium(\n", " tokamak=tokamak,\n", " Rmin=0.1, Rmax=2.0, # radial range\n", " Zmin=-2.2, Zmax=2.2, # vertical range\n", " nx=65, # number of grid points in the radial direction (needs to be of the form (2**n + 1) with n being an integer)\n", " ny=129, # number of grid points in the vertical direction (needs to be of the form (2**n + 1) with n being an integer)\n", " # psi=plasma_psi\n", ") \n", "\n", "# initialise profile object\n", "from freegsnke.jtor_update import ConstrainPaxisIp\n", "profiles = ConstrainPaxisIp(\n", " eq=eq,\n", " paxis=8.1e3,\n", " Ip=6.2e5,\n", " fvac=0.5,\n", " alpha_m=1.8,\n", " alpha_n=1.2\n", ")\n", "\n", "# initialise solver\n", "from freegsnke import GSstaticsolver\n", "GSStaticSolver = GSstaticsolver.NKGSsolver(eq) \n", "\n", "# set coil currents\n", "import pickle\n", "with open('data/simple_diverted_currents_PaxisIp.pk', 'rb') as f:\n", " current_values = pickle.load(f)\n", "\n", "for key in current_values.keys():\n", " eq.tokamak[key].current = current_values[key]\n", "eq.tokamak[\"P6\"].current += 100\n", "\n", "# carry out forward solve\n", "GSStaticSolver.solve(eq=eq, \n", " profiles=profiles, \n", " constrain=None, \n", " target_relative_tolerance=1e-9)\n", "\n", "# plot the resulting equilbrium\n", "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n", "ax1.grid(True, which='both')\n", "eq.plot(axis=ax1, show=False)\n", "eq.tokamak.plot(axis=ax1, show=False)\n", "ax1.set_xlim(0.1, 2.15)\n", "ax1.set_ylim(-2.25, 2.25)\n", "plt.tight_layout()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Define the plasma descriptors\n", "\n", "Here we define the \"plasma descriptors\" (i.e. the shape parameters) that we want to track over time. These descriptors will be evolved according to the following equation:\n", "\n", "$$ \n", " \\vec{s}(t) = \\vec{s}(0) + \\frac{\\partial \\vec{s}}{\\partial \\vec{I}_e} (\\vec{I}_e(t) - \\vec{I}_e(0)) + \\frac{\\partial \\vec{s}}{\\partial \\vec{\\theta}} (\\vec{\\theta}(t) - \\vec{\\theta}(0))\n", "$$\n", "\n", "where\n", " - $\\vec{s}(t)$ are the plasma descriptors (defined in the following cells) approximated via the linearisation, and $\\vec{s}(0)$ are the values defined by the initial equilibrium configuration. \n", " - $\\vec{I}_e(t) = (\\vec{I}_m(t), I_p(t))$ is a vector of the currents in the metals and total plasma current at time $t$. \n", " - $\\vec{\\theta}(t)$ are the (plasma current density) profile parameters at time $t$. \n", " - $\\frac{\\partial \\vec{s}}{\\partial \\vec{I}_e}$ and $\\frac{\\partial \\vec{s}}{\\partial \\vec{\\theta}}$ are the Jacobians calculated wrt the initial equilibrium. \n", "\n", "The currents come from solving the circuit equations the same way as before while the plasma profile parameters are inputs provided by the user. \n", "\n", "This enables a rapid (approximate) evolution of these descriptors over short periods of time - noting that the Jacobians lose accuracy as they diverge from the equilibrium from which they were constructed." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The plasma descriptors are defined via a function that must take a FreeGSNKE equilibrium object as its sole input. \n", "\n", "For example, below we define a function that will return an array containing: the average $Z$ position of the $J_{\\phi}$ map; the radial and vertical position of the lower X-point; and the inboard midplane radius. \n", "\n", "This can be customised to add any descriptors of interest to you. " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# define the descriptors function (it should return an array of values and take in an eq object)\n", "def plasma_descriptors(eq):\n", "\n", " # find lower X-point\n", " # define a \"box\" in which to search for the lower X-point\n", " XPT_BOX = [[0.33, -0.88], [0.95, -1.38]]\n", "\n", " # mask those points\n", " xpt_mask = (\n", " (eq.xpt[:, 0] >= XPT_BOX[0][0])\n", " & (eq.xpt[:, 0] <= XPT_BOX[1][0])\n", " & (eq.xpt[:, 1] <= XPT_BOX[0][1])\n", " & (eq.xpt[:, 1] >= XPT_BOX[1][1])\n", " )\n", " xpts = eq.xpt[xpt_mask, 0:2].squeeze()\n", " if xpts.ndim > 1 and xpts.shape[0] > 1:\n", " opt = eq.opt[0, 0:2]\n", " dists = np.linalg.norm(xpts - opt, axis=1)\n", " idx = np.argmin(dists) # index of closest point\n", " Rx, Zx = xpts[idx, :]\n", " else:\n", " Rx, Zx = xpts\n", "\n", " # find avg. Z position of jtor\n", " Zcurrent = eq.Zcurrent()\n", "\n", " # inboard midplane radius\n", " Rin = eq.innerOuterSeparatrix()[0]\n", "\n", " return np.array([Zcurrent, Rx, Zx, Rin])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Time evolution\n", "\n", "Having defined the plasma descriptors, we can now instantiate the evolutive solver object. By including the `plasma_descriptors` function as an argument, the relevant Jacobians will be calculated to enable this evolution.\n", "\n", "The new **`n_linearization_workers`** argument controls the number of independent worker processes used to construct both the current-response and profile-parameter-response Jacobian columns. `1` retains serial construction; a larger value parallelises the initial linearisation and any later relinearisations. It does not limit the inner threading of ordinary time-step GS solves." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from freegsnke import nonlinear_solve\n", "\n", "stepping = nonlinear_solve.nl_solver(\n", " eq=eq, \n", " profiles=profiles, \n", " GSStaticSolver=GSStaticSolver,\n", " full_timestep=5e-4, \n", " plasma_resistivity=1e-6,\n", " max_mode_frequency=10**2.5,\n", " n_linearization_workers=1, # set >1 to parallelise Jacobian construction\n", " plasma_descriptor_function=plasma_descriptors\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Carry out time stepping\n", "\n", "As we did in the previous notebook, we define the number of steps we wish to take (based on the timestep chosen above) and create some storage lists. " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# number of time steps to simulate\n", "max_count = 50\n", "\n", "# initialising some variables for iteration and logging\n", "counter = 0\n", "t = 0\n", "\n", "# initialise object\n", "stepping.initialize_from_ICs(eq, profiles)\n", "\n", "# storage\n", "history_times = [t]\n", "history_currents = [stepping.currents_vec]\n", "history_plasma_descriptors = [plasma_descriptors(eq)]" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# constant voltages (i.e. no control or external drive)\n", "voltages = (stepping.vessel_currents_vec*stepping.evol_metal_curr.coil_resist)[:stepping.evol_metal_curr.n_active_coils] \n", "\n", "# # time-dependent plasma current density profile parameters\n", "# alpha_m = np.tile(profiles.alpha_m, max_count+1)\n", "# alpha_m -= (0.1 * np.sin(0.05 * np.pi * np.arange(max_count+1))) # we add some perturbation\n", "\n", "# alpha_n = np.tile(profiles.alpha_n, max_count+1)\n", "# alpha_n += (0.1 * np.sin(0.1 * np.pi * np.arange(max_count+1))) # we add some perturbation\n", "\n", "# paxis = np.tile(profiles.paxis, max_count+1)\n", "# paxis += (0.1 * np.sin(0.01 * np.pi * np.arange(max_count+1))) # we add some perturbation" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We can now call the timestepper with both the `linear_only=True` and `no_GS=True` options. This ensures that a linear solve is carried out **without** solving GS at each timestep. Note that you will **not** be able to extract all of the usual equilibrium information at each timestep as the internal `eq` objects no longer contain valid GS solutions! Switch back to `no_GS=False` if that's what you need. \n", "\n", "This means that in this mode of evolution, only the the metal currents, the plasma current, and the plasma descriptors are evolved over time." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# loop over time steps\n", "while counter < max_count:\n", " print(f'Step: {counter}/{max_count-1}')\n", " print(f'--- t = {t:.2e}')\n", "\n", " # carry out the time step (feed in the voltages and profile parameters)\n", " stepping.nlstepper(\n", " active_voltage_vec=voltages,\n", " linear_only=True,\n", " no_GS=True,\n", " verbose=False,\n", " )\n", "\n", " # store time-advanced currents and plasma descriptors\n", " history_currents.append(stepping.currents_vec)\n", " history_plasma_descriptors.append(stepping.plasma_descriptors_vec)\n", "\n", " t += stepping.dt_step\n", " history_times.append(t)\n", " counter += 1\n", "\n", "# transform lists to arrays\n", "history_currents = np.array(history_currents)\n", "history_times = np.array(history_times)\n", "history_plasma_descriptors = np.array(history_plasma_descriptors)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Let us carry out the same evolution, this time with the linear solver actually solving GS at each step (to compare the difference). Note that we change the way we extract the plasma descriptors as `no_GS=False`!" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# re-initialising some variables for iteration and logging\n", "counter = 0\n", "t = 0\n", "\n", "# re-initialise object\n", "stepping.initialize_from_ICs(eq, profiles)\n", "\n", "# storage\n", "history_times_with_GS = [t]\n", "history_currents_with_GS = [stepping.currents_vec]\n", "history_plasma_descriptors_with_GS = [plasma_descriptors(eq)]\n", "\n", "# loop over the time steps\n", "while counter < max_count:\n", " print(f'Step: {counter}/{max_count-1}')\n", " print(f'--- t = {t:.2e}')\n", " \n", " # carry out the time step (feed in the voltages and profile parameters)\n", " stepping.nlstepper(\n", " active_voltage_vec=voltages,\n", " linear_only=True,\n", " no_GS=False,\n", " verbose=False,\n", " )\n", "\n", " # store time-advanced currents and plasma descriptors\n", " history_currents_with_GS.append(stepping.currents_vec)\n", " history_plasma_descriptors_with_GS.append(plasma_descriptors(stepping.eq1)) # <-- changed!\n", "\n", " t += stepping.dt_step\n", " history_times_with_GS.append(t)\n", " counter += 1\n", "\n", "# transform lists to arrays\n", "history_currents_with_GS = np.array(history_currents_with_GS)\n", "history_times_with_GS = np.array(history_times_with_GS)\n", "history_plasma_descriptors_with_GS = np.array(history_plasma_descriptors_with_GS)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now we can compare the difference between the linear evolution with and without GS solves at each step. We see that for some parameters the evolution diverges very quickly with solving GS!" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, ax = plt.subplots(4,1, figsize=(10, 12), dpi=80)\n", "ax = ax.flatten()\n", "\n", "labels = [\"Zcurrent [m]\", \"Rx [m]\", \"Zx [m]\", \"Rin [m]\"]\n", "\n", "for i in range(4):\n", " ax[i].plot(history_times_with_GS, history_plasma_descriptors_with_GS[:, i], 'k+', label=\"Linear (with GS)\")\n", " ax[i].plot(history_times, history_plasma_descriptors[:, i], 'm3', label=\"Linear (without GS)\")\n", " ax[i].set_xlabel(\"Time (s)\")\n", " ax[i].set_ylabel(labels[i])\n", " ax[i].legend()\n", " ax[i].grid()\n", "\n", "fig.tight_layout()\n" ] } ], "metadata": { "kernelspec": { "display_name": ".venv", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.13" } }, "nbformat": 4, "nbformat_minor": 4 }